#!/usr/bin/env python

import json
from nacl.bindings.crypto_scalarmult import crypto_scalarmult
from nacl.signing import SigningKey, VerifyKey
from nacl.bindings.crypto_aead import crypto_aead_chacha20poly1305_ietf_encrypt, crypto_aead_chacha20poly1305_ietf_decrypt
import hashlib
from nacl.encoding import Base64Encoder
import nacl.utils
import sys
import os


KEY_SIZE = 32
NONCE_SIZE = 12
SIGN_SIZE = 64


def create_encrypted_packet(keys, data, nonce=None):
	sender_private_key = Base64Encoder.decode(keys["sender_private_key"])
	sender_sign_private_key = SigningKey(keys["sender_sign_private_key"], Base64Encoder)
	receiver_public_key = Base64Encoder.decode(keys["receiver_public_key"])
	receiver_sign_public_key = VerifyKey(keys["receiver_sign_public_key"], Base64Encoder)

	ss = crypto_scalarmult(sender_private_key, receiver_public_key)

	key = nacl.utils.random(KEY_SIZE)
	if nonce is None:
		nonce = os.urandom(NONCE_SIZE)

	hash = hashlib.sha256()
	hash.update(key)
	hash.update(data)
	hash = hash.digest()

	sign = sender_sign_private_key.sign(hash).signature

	ret = bytes(bytearray(ord(ss[i]) ^ ord(key[i]) for i in range(KEY_SIZE)))
	ret += nonce
	ret += crypto_aead_chacha20poly1305_ietf_encrypt(sign + data, None, nonce, key)
	return ret


def verify_encrypted_packet(keys, data):
	sender_private_key = Base64Encoder.decode(keys["sender_private_key"])
	sender_sign_private_key = SigningKey(keys["sender_sign_private_key"], Base64Encoder)
	receiver_public_key = Base64Encoder.decode(keys["receiver_public_key"])
	receiver_sign_public_key = VerifyKey(keys["receiver_sign_public_key"], Base64Encoder)

	ss = crypto_scalarmult(sender_private_key, receiver_public_key)

	key = bytes(bytearray(ord(ss[i]) ^ ord(data[i]) for i in range(KEY_SIZE)))

	dec_data = crypto_aead_chacha20poly1305_ietf_decrypt(data[KEY_SIZE+NONCE_SIZE:], None, data[KEY_SIZE:KEY_SIZE+NONCE_SIZE], key)
	sign = dec_data[:SIGN_SIZE]
	dec_data = dec_data[SIGN_SIZE:]

	hash = hashlib.sha256()
	hash.update(key)
	hash.update(dec_data)
	hash = hash.digest()

	receiver_sign_public_key.verify(hash, sign)

	return dec_data


def load_key_file(filename):
	with open(filename, 'r') as fp:
		return json.load(fp)


def usage():
	print("Usage: %s (-d | -e) key_file input_file output_file" % sys.argv[0])
	sys.exit(1)


def main():
	if len(sys.argv) < 5:
		usage()
	if sys.argv[1] == "-e":
		keys = load_key_file(sys.argv[2])
		with open(sys.argv[3], 'rb') as fp:
			data = fp.read()
		data = create_encrypted_packet(keys, data)
		with open(sys.argv[4], 'wb') as fp:
			fp.write(data)
	elif sys.argv[1] == "-d":
		keys = load_key_file(sys.argv[2])
		with open(sys.argv[3], 'rb') as fp:
			data = fp.read()
		data = verify_encrypted_packet(keys, data)
		with open(sys.argv[4], 'wb') as fp:
			fp.write(data)
	else:
		usage()


if __name__ == '__main__':
	main()
